import os
import torch.utils.data as data

from PIL import Image
import numpy as np

def voc_cmap(N=256, normalized=False):
    def bitget(byteval, idx):
        return ((byteval & (1 << idx)) != 0)

    dtype = 'float32' if normalized else 'uint8'
    cmap = np.zeros((N, 3), dtype=dtype)
    for i in range(N):
        r = g = b = 0
        c = i
        for j in range(8):
            r = r | (bitget(c, 0) << 7-j)
            g = g | (bitget(c, 1) << 7-j)
            b = b | (bitget(c, 2) << 7-j)
            c = c >> 3

        cmap[i] = np.array([r, g, b])

    cmap = cmap/255 if normalized else cmap
    return cmap


class SelfDataset(data.Dataset):

    cmap = voc_cmap()

    def __init__(self, imgs_dir, anno_dir, split='train', transform=None):
        self.imgs_dir = imgs_dir
        self.anno_dir = anno_dir
        self.transform = transform
        if not os.path.isdir(self.imgs_dir) or not os.path.isdir(self.anno_dir):
            raise RuntimeError('Dataset not found or corrupted.')
        with open('{}.txt'.format(split), 'r') as f:
            self.id_list = list(map(lambda x: str(int(x.strip('\n'))), f.readlines()))
        self.images = []
        self.gts = []
        for i_id in self.id_list:
            i_id = i_id.replace('\r', '').replace('\n', '').replace(' ', '')
            self.images.append(os.path.join(self.imgs_dir, i_id + '.jpg'))
            self.gts.append(os.path.join(self.anno_dir, i_id + '.png'))

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            tuple: (image, target) where target is the image segmentation.
        """
        img = Image.open(self.images[index]).convert('RGB')
        target = Image.open(self.gts[index])
        W, H = img.size
        if self.transform is not None:
            img, target = self.transform(img, target)
        return self.id_list[index], img, target, W, H

    def __len__(self):
        return len(self.images)

    @classmethod
    def decode_target(cls, mask):
        """decode semantic mask to RGB image"""
        return cls.cmap[mask]


class TestDataset(data.Dataset):
    cmap = voc_cmap()

    def __init__(self, imgs_dir, transform=None):
        self.imgs_dir = imgs_dir
        self.transform = transform
        with open('{}.txt'.format('test'), 'r') as f:
            self.id_list = list(map(lambda x: str(int(x.strip('\n'))), f.readlines()))
        self.images = []
        for i_id in self.id_list:
            i_id = i_id.replace('\r', '').replace('\n', '').replace(' ', '')
            self.images.append(os.path.join(self.imgs_dir, i_id + '.jpg'))

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            tuple: (image, target) where target is the image segmentation.
        """
        img = Image.open(self.images[index]).convert('RGB')
        W, H = img.size
        if self.transform is not None:
            img, tmp = self.transform(img, img)
        return self.id_list[index], img, W, H

    def __len__(self):
        return len(self.images)

    @classmethod
    def decode_target(cls, mask):
        """decode semantic mask to RGB image"""
        return cls.cmap[mask]
